
#加载MNIST数据集
from tensorflow import keras
import matplotlib.pyplot as plt

# 加载数据集（MNIST）
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

fig = plt.figure()

# 可视化一些数字
for i in range(9):
    plt.subplot(3,3,i+1)
    plt.tight_layout()
    plt.imshow(x_train[i], interpolation='none')
    plt.title("Digit: {}".format(y_train[i]))

    # 不显示x和y轴刻度
    plt.xticks([])
    plt.yticks([])

#对图像数据进行归一化，并扩展数据集的维度以准备训练
import tensorflow as tf
import numpy as np

# 归一化图像
x_train, x_test = x_train / 255.0, x_test / 255.0

# 扩展维度以适应CNN
x_train = np.expand_dims(x_train, axis=-1)
x_test = np.expand_dims(x_test, axis=-1)

# 将标签转换为分类格式（one-hot编码）
y_train = keras.utils.to_categorical(y_train, 10)
y_test = keras.utils.to_categorical(y_test, 10)


#定义一个教师模型——一个包含多个层的CNN
# 教师模型
teacher_model = keras.Sequential([
    keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Conv2D(64, (3, 3), activation='relu'),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Flatten(),
    keras.layers.Dense(128, activation='relu'),
    keras.layers.Dense(10)  # 不使用softmax，输出原始logits以用于蒸馏
])

# 注意，学生模型的最后一层有10个单元（每个数字一个），但没有使用softmax激活函数。它输出原始logits，这在模型蒸馏过程中非常重要，因为softmax会在蒸馏过程中应用于计算教师模型和学生模型之间的Kullback-Leibler（KL）散度。
# 定义好教师神经网络后，使用compile()方法配置优化器、损失函数和评估指标：
teacher_model.compile(
    optimizer = 'adam',
    loss = tf.keras.losses.CategoricalCrossentropy(from_logits = True),
    metrics = ['accuracy']
)
# 训练教师模型
teacher_model.fit(x_train, y_train, 
                  epochs = 5, 
                  batch_size = 64, 
                  validation_data = (x_test, y_test))


# 定义学生模型
# 教师模型训练完成后，接下来定义学生模型。学生模型是一个比教师模型更简单的架构，层数更少：

# 学生模型
student_model = keras.Sequential([
    keras.layers.Conv2D(16, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    keras.layers.MaxPooling2D((2, 2)),
    keras.layers.Flatten(),
    keras.layers.Dense(64, activation='relu'),
    keras.layers.Dense(10)  # 不使用softmax，输出原始logits以用于蒸馏
])


# 定义蒸馏损失函数
# 接下来，定义蒸馏函数，该函数将使用教师模型的预测和学生模型的预测来计算蒸馏损失。该函数将：

# 使用教师模型生成输入批次的软目标（概率）。

# 使用学生模型的预测计算其软概率。

# 计算教师模型和学生模型软概率之间的Kullback-Leibler（KL）散度。

# 返回蒸馏损失。

# 软概率指的是多个可能结果的概率分布，而不是分配一个硬标签。例如，在一个预测电子邮件是否为垃圾邮件的机器学习模型中，模型可能会输出一个概率，如0.85表示垃圾邮件，0.15表示非垃圾邮件。这意味着模型有85%的置信度认为邮件是垃圾邮件，但仍认为有15%的可能性不是，从而允许更好的决策和阈值调整。

# 软概率使用softmax函数计算，并通过温度参数控制。在知识蒸馏中，教师模型的软概率帮助学生模型学习关于类别之间关系的更丰富信息，从而提高泛化能力和性能。

# 以下是distillation_loss()函数的定义：
def distillation_loss(y_true, y_pred, x_batch, teacher_model, temperature=5):
    """
    使用KL散度计算蒸馏损失。
    """
    # 计算当前批次的教师logits
    teacher_logits = teacher_model(x_batch, training=False)

    # 将logits转换为软概率
    teacher_probs = tf.nn.softmax(teacher_logits / temperature)
    student_probs = tf.nn.softmax(y_pred / temperature)

    # KL散度损失（教师和学生分布之间的差异）
    return tf.reduce_mean(tf.keras.losses.KLDivergence()(teacher_probs, student_probs))

# Kullback-Leibler（KL）散度，也称为相对熵，是衡量一个概率分布与另一个参考概率分布之间差异的指标。
# 使用知识蒸馏训练学生模型
# 现在，你可以使用知识蒸馏来训练学生模型。首先，定义train_step()函数：
optimizer = tf.keras.optimizers.Adam()

@tf.function
def train_step(x_batch, y_batch, student_model, teacher_model):
    with tf.GradientTape() as tape:
        # 获取学生模型的预测
        student_preds = student_model(x_batch, training=True)

        # 计算蒸馏损失（显式传递teacher_model）
        loss = distillation_loss(y_batch, student_preds, x_batch, teacher_model, temperature=5)

    # 计算梯度
    gradients = tape.gradient(loss, student_model.trainable_variables)

    # 应用梯度——训练学生模型
    optimizer.apply_gradients(zip(gradients, student_model.trainable_variables))

    return loss


# 该函数执行单次训练步骤：

# 计算学生模型的预测。

# 使用教师模型的预测计算蒸馏损失。

# 计算梯度并更新学生模型的权重。

# 为了训练学生模型，你创建一个训练循环来迭代数据集，在每一步更新学生模型的权重，并在每个epoch结束时打印损失以监控进度：

# 训练循环
epochs = 5
batch_size = 32

# 准备数据集批次
train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)

for epoch in range(epochs):
    total_loss = 0
    num_batches = 0

    for x_batch, y_batch in train_dataset:
        loss = train_step(x_batch, y_batch, student_model, teacher_model)
        total_loss += loss.numpy()
        num_batches += 1

    avg_loss = total_loss / num_batches
    print(f"Epoch {epoch+1}, Loss: {avg_loss:.4f}")

print("学生模型训练完成！")


# 评估学生模型
# 学生模型训练完成后，你可以使用测试集（x_test和y_test）评估其性能：
student_model.compile(
    optimizer='adam',
    loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
    metrics=['accuracy']
)

student_acc = student_model.evaluate(x_test, y_test, verbose=0)[1]
print(f"学生模型准确率: {student_acc:.4f}")

# 使用教师模型和学生模型进行预测
# 现在，你可以使用教师模型和学生模型进行一些预测，看看两者是否能够准确预测MNIST测试数据集中的数字：
import numpy as np

_, (x_test, y_test) = keras.datasets.mnist.load_data()

for index in range(5):    
    plt.figure(figsize=(2, 2))
    plt.imshow(x_test[index], interpolation='none')
    plt.title("Digit: {}".format(y_test[index]))
    # 不显示x和y轴刻度
    plt.xticks([])
    plt.yticks([])
    plt.show()
    
    # 现在可以进行预测
    x = x_test[index].reshape(1,28,28,1)
    
    predictions = teacher_model.predict(x)
    print(predictions)
    print("教师模型预测值: ", np.argmax(predictions, axis=-1))
    
    predictions = student_model.predict(x)
    print(predictions)
    print("学生模型预测值: ", np.argmax(predictions, axis=-1))

# 你会发现学生模型的表现与教师模型一样好。

# 总结
# 在本文中，我们探讨了模型蒸馏的概念，这是一种让更小、更简单的学生模型模仿更大、更复杂的教师模型性能的技术。我们通过MNIST数据集训练了一个教师模型，然后应用蒸馏技术训练了一个学生模型。学生模型通过更少的层和更低的复杂度，成功地模仿了教师模型的性能，同时减少了计算资源的需求。